/* This is the code which stores a VLMM as a n-ary Tree capable of searching through the Tree
Basic operations required are
1. Build a Tree
2. Search a Tree
3. Query probability of a context
4. Build a tree in real time 
*/

#define BACKOFF_DISCOUNT (0.4)
typedef struct Treenode* pstNode;
typedef struct probData* pstData;
struct probData 
{
	int symb;
    float prob;
    int myCount; 
    int childCount;
    char Id;
};

struct Treenode
{
	pstNode parent;
	pstNode child[NUM_GROOVES];
	pstData data;
	int level;
};

typedef struct srch_struct{
	int found;
	pstNode ptrSeq[MAX_DEPTH];
	int seqLen;
} srch;

typedef struct query_struct{
	int found;
	float prob;
} prob;


pstNode createTree(void);
pstNode addNode(pstNode parentNode, pstNode node, int ChildIndex);
pstData createNewData(void);
pstNode createNode(pstData data);
pstNode growPST(pstNode PST, int* seq, int Len, int MaxDepth);
int getParentSymbol(pstNode);
pstNode addNodeToPST(pstNode pst, pstNode node, pstNode parent, int ChildIndex);
srch searchTree(pstNode pst, int* seq, int L, int R);
prob getCondProbBackoff(pstNode PST, int* seq, int L, int R, int MaxDepth);
void getPredSymbDist(pstNode PST, int* cont, int L, int R, int NumSym, float* probab);
pstNode buildPSTRealTime(pstNode PST, int* seq, int L, int R);
int findmax(int a, int b);

pstData createNewData(void){
	pstData data;
	data = (pstData) malloc(sizeof(struct probData));
	data->symb = 0;
	data->prob = 0;
	data->myCount = 0; 
	data->childCount = 0;
	data->Id = 'a';
	return data;
}

pstNode createNode(pstData data){
	pstNode t;
	int i;
	t=(pstNode) malloc(sizeof(struct Treenode));
	t->data = data;
	assert(t!=NULL);
	t->parent = NULL;
	t->level = -1;
	for(i=0;i<NUM_GROOVES;i++){
		t->child[i] = NULL;
	}
	return t;
}

pstNode createTree(void){
	pstNode t;
	pstData data;
	data = createNewData();
	data->symb = 0;
	t = createNode(data);
	t->parent = NULL;
	t->level = 0;
	return t;
}
void printNode(pstNode pst){
	int par, k;
	if(pst == NULL)
		printf("Empty Tree!\n");
	else {
		printf("Node Info:\n");
		printf("Node level: %d\nNode data: %d, %f, %d, %d, %c\n",pst->level, pst->data->symb, pst->data->prob, pst->data->myCount, pst->data->childCount, pst->data->Id);
		par = getParentSymbol(pst);
		printf("Node parent symbol: %d\n",par);
		printf("Child Status: \n");
		for(k=0;k<NUM_GROOVES;k++){
			//printf("Child %d:%d\n",k,(int)((int)(pst->child[k]) > 0));
		}
	}
}

int getParentSymbol(pstNode pst){
	if(pst->parent != NULL) {
			return (pst->parent->data->symb);
		}
		else {
			return -1;
		}
}

pstNode addNodeToPST(pstNode pst, pstNode node, pstNode parent, int ChildIndex){
	if(parent->child[ChildIndex] != NULL){
//		printf("Replacing the existing child!");
	}
	node->parent = parent;
	node->level = parent->level + 1;
	parent->child[ChildIndex] = node;
	return pst;
}

srch searchTree(pstNode pst, int* seq, int L, int R){
	int currlen;
	int k,j;
	pstNode t = pst;
	srch srchResult;
	srchResult.found = 0;
	j = 0;
	srchResult.ptrSeq[j++] = pst;
	currlen = R-L+1;
	for(k=L;k<=R;k++){
		srchResult.ptrSeq[j] = srchResult.ptrSeq[j-1]->child[seq[k]-1];
		if(srchResult.ptrSeq[j] == NULL)
			break;
		else
			j++;
	}
	if(k == R+1){
		srchResult.found = 1;
		srchResult.ptrSeq[j] = NULL;
	}
	srchResult.seqLen = j;
	return srchResult;
}

pstNode growPST(pstNode PST, int* seq, int Len, int MaxDepth){
	pstData data;
	pstNode node;
	int d,k;
	srch sr;
	for(d=0;d<MaxDepth;d++){
		for(k=d;k<Len;k++){
			sr = searchTree(PST,seq,k-d,k);
			if(!sr.found){
				data = createNewData();
				data->symb = seq[k];
				data->myCount = 1;
				node = createNode(data);
				PST = addNodeToPST(PST, node, sr.ptrSeq[sr.seqLen-1], (data->symb)-1);
				sr.ptrSeq[sr.seqLen-1]->data->childCount++;
			}
			else{
				sr.ptrSeq[sr.seqLen-1]->data->myCount++;
				sr.ptrSeq[sr.seqLen-2]->data->childCount++;
			}
		}
	}
	return PST;
}

prob getCondProbBackoff(pstNode PST, int* seq, int L, int R, int MaxDepth){
	int seqLen = R-L+1;
	prob var;
	srch sr;
	int k;
	float discount = 1.0;
	var.found = 0;
	var.prob = 0.0;
	if(R-L+1 > MaxDepth){
		printf("Sequence length greater than VLMM Order\n");
		return var;
	}
	for(k=0;k<seqLen;k++){
		sr = searchTree(PST,seq,k+L,R);
		if(sr.found){
			var.prob = discount * (float)(sr.ptrSeq[sr.seqLen-1]->data->myCount)/((float)(sr.ptrSeq[sr.seqLen-2]->data->childCount));
			var.found = 1;
			break;
		}
		discount = discount * BACKOFF_DISCOUNT;
	}
	return var;
}

pstNode buildPSTRealTime(pstNode PST, int* seq, int L, int R){
	pstData data;
	pstNode node;
	int k;
	srch sr;
	for(k=0;k<(R-L+1);k++){			// Update all n-grams
		sr = searchTree(PST,seq,R-k,R);
		if (!sr.found){
			data = createNewData();
			data->symb = seq[R];
			data->myCount = 1;
			node = createNode(data);
			PST = addNodeToPST(PST, node, sr.ptrSeq[sr.seqLen-1], (data->symb)-1);
			sr.ptrSeq[sr.seqLen-1]->data->childCount++;
		}
		else{
			sr.ptrSeq[sr.seqLen-1]->data->myCount++;
			sr.ptrSeq[sr.seqLen-2]->data->childCount++;
		}
	}
	return PST;
}

int findmax(int a, int b){
	if (a>b)
		return a;
	else
		return b;
}

void getPredSymbDist(pstNode PST, int* cont, int L, int R, int NumSym, float* probab){
	prob pr;
	float sumProb = 0.0;
	int context[MAX_DEPTH] = {0};
	int k = 0;
	int symbIndex;
	for(k=0;k<(R-L+1);k++){
		context[k] = cont[k+L];
	}
	symbIndex = k;
	for(k=0;k<NUM_GROOVES;k++){
		context[symbIndex] = k+1;		// Next symbol 
		pr = getCondProbBackoff(PST, context, 0, R-L+1, R-L+2);
		probab[k] = pr.prob;
		sumProb = sumProb + probab[k];
	}
	for(k=0;k<NUM_GROOVES;k++){
		probab[k] = probab[k]/sumProb;		//Normalizing
	}
	return;
}
